from rembg import remove, new_session
from PIL import Image
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 全局模型会话（避免每次请求都加载模型）
try:
    logger.info("正在加载AI模型...")
    session = new_session(model_name='u2net_human_seg')
    logger.info("AI模型加载完成")
except Exception as e:
    logger.error(f"模型加载失败: {str(e)}")
    raise RuntimeError(f"模型加载失败: {str(e)}")


def remove_background(input_img: Image.Image) -> Image.Image:
    """
    移除背景并转为白底

    参数:
    input_img: 输入图像(PIL格式)

    返回:
    白底人像图像(PIL格式)
    """
    try:
        # 移除背景（生成透明背景）
        output_img = remove(input_img, session=session)

        # 创建白色背景
        white_bg = Image.new("RGB", output_img.size, (255, 255, 255))

        # 将透明图像叠加到白色背景上
        if output_img.mode == 'RGBA':
            # 使用alpha通道作为蒙版
            white_bg.paste(output_img, mask=output_img.split()[3])
        else:
            # 如果没有alpha通道，直接粘贴
            white_bg.paste(output_img)

        return white_bg

    except Exception as e:
        logger.error(f"背景移除失败: {str(e)}")
        raise RuntimeError(f"背景移除失败: {str(e)}")